{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import some stuff\n",
    "\n",
    "import os\n",
    "import sys\n",
    "import glob\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import math\n",
    "# from pathlib import Path\n",
    "\n",
    "from evo.core import metrics\n",
    "from evo.core.trajectory import PoseTrajectory3D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download NTUVIRAL Groundtruth,\n",
    "!rm -rf ntuviral_gt\n",
    "!git clone https://github.com/ntu-aris/ntuviral_gt\n",
    "\n",
    "# Set the ground truth path here\n",
    "gtgen_res_path = 'ntuviral_gt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Or set the path to a different path\n",
    "# gtgen_res_path = '/home/tmn/DATASETS/NTUVIRAL/ntuviral_gt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Downdload sample fast-lio2 estimate\n",
    "!rm -rf fastlio2_sample\n",
    "!git clone https://github.com/ntu-aris/fastlio2_sample\n",
    "\n",
    "# Set the path to the logs of your slam estimate\n",
    "slam_est_path = 'fastlio2_sample'\n",
    "\n",
    "# Set the directory where results are exported\n",
    "output_dir = slam_est_path + '/analysis'\n",
    "os.makedirs(output_dir, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Conversion between quat and rotm\n",
    "from evo.core.transformations import quaternion_matrix as quat2rotm_\n",
    "\n",
    "def quat2rotm(q):\n",
    "    return quat2rotm_(q)[0:3, 0:3]\n",
    "\n",
    "# region Customizable for each method ---------------------\n",
    "\n",
    "# Offset from body center to the prism\n",
    "t_B_prism = np.array([-0.293656, -0.012288, -0.273095]).reshape((3,1))\n",
    "\n",
    "# endregion Customizable for each method -----------------\n",
    "\n",
    "algo_name = 'fastlio2'\n",
    "\n",
    "# Minimum completeness to judge ate\n",
    "min_completeness = 90.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Search for the sequence name\n",
    "gndtr_logs = glob.glob(gtgen_res_path + '/**/ground_truth.csv', recursive=True)\n",
    "gndtr_logs = sorted(gndtr_logs)\n",
    "gndtr_df   = pd.DataFrame([str(x) for x in gndtr_logs], columns=['fullpath'])\n",
    "\n",
    "# Check for the sequence\n",
    "def decode_gndtr_sequence_name(x):\n",
    "    sequence = x.split('/')[-2]\n",
    "    return sequence\n",
    "\n",
    "gndtr_df['sequence'] = gndtr_df['fullpath'].apply(lambda x : decode_gndtr_sequence_name(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the ground truth array for each sequence\n",
    "def load_csv(log):\n",
    "    \n",
    "    # Open the log\n",
    "    data = np.loadtxt(log, delimiter=',', skiprows=1)\n",
    "\n",
    "    # Pose stamp\n",
    "    t    = data[:, 2]/1e9\n",
    "    xyz  = data[:, 3:6]\n",
    "    quat = data[:, 6:10]\n",
    "\n",
    "    pose_stamp = np.empty((data.shape[0], 8))\n",
    "    pose_stamp[:, 0]   = t\n",
    "    pose_stamp[:, 1:4] = xyz\n",
    "    pose_stamp[:, 4:8] = quat\n",
    "    \n",
    "    # Create the spline\n",
    "    return pose_stamp\n",
    "\n",
    "gndtr_df['pose_stamped'] = gndtr_df['fullpath'].apply(lambda x : load_csv(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Search for the estimates\n",
    "slam_est_logs = glob.glob(slam_est_path + '/**/odometry.csv', recursive=True)\n",
    "slam_est_logs = sorted(slam_est_logs)\n",
    "est_df = pd.DataFrame([str(x) for x in slam_est_logs if 'result_' in str(x)], columns=['fullpath'])\n",
    "\n",
    "# Decode the sequence\n",
    "def decode_est_sequence_name(x):\n",
    "    dirname = os.path.dirname(x)\n",
    "    seqname = dirname.split('/')[-1].replace('result_', '')\n",
    "    return seqname\n",
    "\n",
    "est_df['sequence'] = est_df['fullpath'].apply(lambda x : decode_est_sequence_name(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getGTMinTime(x):\n",
    "    print(x)\n",
    "    return gndtr_df[gndtr_df['sequence'] == x]['pose_stamped'].iloc[0][0, 0]\n",
    "\n",
    "def getGTMaxTime(x):\n",
    "    print(x)\n",
    "    return gndtr_df[gndtr_df['sequence'] == x]['pose_stamped'].iloc[0][-1, 0]\n",
    "\n",
    "est_df['t_min'] = est_df['sequence'].apply( lambda x : getGTMinTime(x) )\n",
    "est_df['t_max'] = est_df['sequence'].apply( lambda x : getGTMaxTime(x) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the estimate data\n",
    "est_df['data_est'] = est_df['fullpath'].apply(lambda x : np.loadtxt(x, delimiter=',', skiprows=1))\n",
    "\n",
    "def extract_est_data(data, t_min, t_max, path = None):\n",
    "    \n",
    "    # print(t_min, 'to', t_max)\n",
    "    t = data[:, 0]/1.0e9\n",
    "    P = data[:, 3:6]\n",
    "    Q = data[:, [9, 6, 7, 8]]\n",
    "    idx_intime = [ idx for idx in range(0, len(t)) if t[idx] >= t_min and t[idx] <= t_max ]\n",
    "\n",
    "    t = t[idx_intime]\n",
    "    P = P[idx_intime, :]\n",
    "    Q = Q[idx_intime, :]\n",
    "\n",
    "    # Add the pose offset\n",
    "    for n in range(0, len(P)):\n",
    "        P[n, :] = P[n, :] + (quat2rotm(Q[n, :]).dot(t_B_prism)).transpose()\n",
    "\n",
    "    return (t, P, Q)\n",
    "\n",
    "est_df['t_est'], est_df['P_est'], est_df['Q_est'] = zip(*map(extract_est_data, est_df['data_est'], est_df['t_min'], est_df['t_max'], est_df['fullpath']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sample the ground truth pose against the estimate.\n",
    "\n",
    "from warnings import warn\n",
    "from evo.core.trajectory import PoseTrajectory3D\n",
    "from evo.core.transformations import quaternion_slerp as qslerp\n",
    "\n",
    "def lerp(a, b, t):\n",
    "    return (1.0-t)*a + t*b\n",
    "\n",
    "# We need to do linear interpolation, ourselves...\n",
    "def cv_interp(traj, ref_timestamps, *, extrapolate_past_end=False):\n",
    "    \"Compute points along trajectory *est* at *timestamps* or trajectory *ref*.\"\n",
    "    # Accept trajectories \n",
    "    if hasattr(ref_timestamps, 'timestamps'):\n",
    "        ref_timestamps = ref_timestamps.timestamps\n",
    "\n",
    "    ref_timestamps = np.array(ref_timestamps, copy=True)\n",
    "    \n",
    "    est_tran = traj.positions_xyz\n",
    "    est_quat = traj.orientations_quat_wxyz\n",
    "    \n",
    "    # Index of closest next estimated timestamp for each pose in *traj*.\n",
    "    # Must be >= 1 since we look at the previous pose.\n",
    "    i = 1\n",
    "\n",
    "    for j, t_ref in enumerate(ref_timestamps):\n",
    "        while i < traj.num_poses and traj.timestamps[i] <= t_ref:\n",
    "            i += 1\n",
    "        if not extrapolate_past_end and traj.num_poses <= i:\n",
    "            warn('reference trajectory ends after estimated, cut short')\n",
    "            break\n",
    "        t_prev = traj.timestamps[i-1]\n",
    "        t_next = traj.timestamps[i]\n",
    "        td     = (t_ref - t_prev) / (t_next - t_prev)\n",
    "        quat_j = qslerp(est_quat[i-1], est_quat[i], td)\n",
    "        tran_j = lerp(est_tran[i-1], est_tran[i], td)\n",
    "        yield np.r_[t_ref, tran_j, quat_j]\n",
    "\n",
    "def trajectory_interpolation(traj, timestamps):\n",
    "    poses = np.array(list(cv_interp(traj, timestamps)))\n",
    "    times = poses[:, 0]\n",
    "    trans = poses[:, 1:4]\n",
    "    quats = poses[:, 4:8]\n",
    "    return PoseTrajectory3D(trans, quats, times, meta=traj.meta)\n",
    "\n",
    "# Make the associable trajectory\n",
    "def make_traj(x):\n",
    "    \n",
    "    t_gt = gndtr_df[gndtr_df['sequence'] == x['sequence']]['pose_stamped'].iloc[0][:, 0]    \n",
    "    valid = []\n",
    "    for n in range(0, len(x['t_est'])):\n",
    "        sample_time = x['t_est'][n]\n",
    "        min_diff = np.min(np.abs(t_gt - sample_time))\n",
    "        if min_diff < 0.05:\n",
    "            valid.append(n)\n",
    "    \n",
    "    t_valid = x['t_est'][valid]\n",
    "    P_valid = x['P_est'][valid, :]\n",
    "    Q_valid = x['Q_est'][valid, :]\n",
    "    \n",
    "    return PoseTrajectory3D(P_valid, Q_valid, t_valid)\n",
    "                            \n",
    "est_df['traj_est'] = est_df.apply(lambda x : make_traj(x), axis=1)\n",
    "\n",
    "def makeGTTraj(x):\n",
    "    pose_stamped = gndtr_df[gndtr_df['sequence'] == x['sequence']]['pose_stamped'].iloc[0]\n",
    "    gtr_traj = PoseTrajectory3D(pose_stamped[:, 1:4], pose_stamped[:, 4:8], pose_stamped[:, 0])\n",
    "    return trajectory_interpolation(gtr_traj, x['traj_est'].timestamps)\n",
    "\n",
    "est_df['traj_gtrest'] = est_df.apply(lambda x : makeGTTraj(x), axis=1)\n",
    "\n",
    "est_df['P_gtrest'] = est_df['traj_gtrest'].apply(lambda x : x.positions_xyz)\n",
    "est_df['Q_gtrest'] = est_df['traj_gtrest'].apply(lambda x : x.orientations_quat_wxyz[:, [3, 0, 1, 2]])\n",
    "est_df['t_gtrest'] = est_df['traj_gtrest'].apply(lambda x : x.timestamps)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Align the traj\n",
    "for i in est_df.index:\n",
    "    est_df.at[i, 'traj_est'].align(est_df.at[i, 'traj_gtrest'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate the ATE\n",
    "def calculate_metric(traj_gtr, traj_est):\n",
    "    metric = metrics.APE(pose_relation=metrics.PoseRelation.translation_part)\n",
    "    metric.process_data((traj_gtr, traj_est))\n",
    "    return float(metric.get_result(ref_name='reference', est_name='estimate').stats['rmse'])\n",
    "\n",
    "est_df['ate_est'] = est_df.apply(lambda x : calculate_metric(x['traj_gtrest'], x['traj_est']), axis=1)\n",
    "est_df['completeness_est'] = est_df.apply(lambda x : round((x['t_est'][-1] -  x['t_min']) / (x['t_max'] - x['t_min'])*100, 0), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extracting the data to generate table of all sequences\n",
    "ate_df = est_df[['sequence', 'ate_est', 'completeness_est', 'fullpath']]\n",
    "ate_df['algo'] = algo_name\n",
    "\n",
    "def check_if_complete(x):\n",
    "    if x['ate_est'] < 20 and x['completeness_est'] < min_completeness:\n",
    "        return float('inf')\n",
    "    else:\n",
    "        return x['ate_est']\n",
    "\n",
    "ate_df['ate_est'] = ate_df.apply(lambda x : check_if_complete(x), axis=1)\n",
    "\n",
    "# ate_df.to_pickle(output_dir + '/ate_df.pkl')\n",
    "\n",
    "def make_ate_txt(odom):\n",
    "        return f'{round(odom, 3):.03f}'\n",
    "\n",
    "ate_df['report'] = ate_df.apply(lambda x : make_ate_txt(x['ate_est']), axis=1)\n",
    "\n",
    "ate_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Making a latex export\n",
    "latex_df = ate_df[['sequence', 'report', 'completeness_est']]\n",
    "\n",
    "# Writing the report\n",
    "print(latex_df.to_latex(index_names=['#'], header=['Seq', 'ATE', '%']))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py3810",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
